"""
AI任务调度器
负责AI任务的队列管理、优先级调度和负载均衡
"""
import asyncio
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Callable, Any
from uuid import uuid4
from collections import deque
import heapq

from schemas.gpu_resource import (
    AITask, TaskStatus, TaskPriority, GPUResourceRequest,
    GPUResourceAllocation, TaskScheduleRequest, TaskScheduleResponse
)
from services.ai_infrastructure.gpu_resource_manager import gpu_resource_manager

logger = logging.getLogger(__name__)


class TaskQueue:
    """任务队列"""
    
    def __init__(self):
        self.tasks: List[AITask] = []
        self.priority_map = {
            TaskPriority.CRITICAL: 0,
            TaskPriority.HIGH: 1,
            TaskPriority.NORMAL: 2,
            TaskPriority.LOW: 3
        }
    
    def add_task(self, task: AITask):
        """添加任务到队列"""
        # 使用优先级和创建时间作为排序键
        priority_value = self.priority_map[task.priority]
        timestamp = task.created_at.timestamp()
        
        # 使用堆来维护优先级队列
        heapq.heappush(self.tasks, (priority_value, timestamp, task))
    
    def get_next_task(self) -> Optional[AITask]:
        """获取下一个任务"""
        if not self.tasks:
            return None
        
        _, _, task = heapq.heappop(self.tasks)
        return task
    
    def remove_task(self, task_id: str) -> bool:
        """移除指定任务"""
        for i, (_, _, task) in enumerate(self.tasks):
            if task.task_id == task_id:
                del self.tasks[i]
                heapq.heapify(self.tasks)  # 重新构建堆
                return True
        return False
    
    def get_task_position(self, task_id: str) -> Optional[int]:
        """获取任务在队列中的位置"""
        for i, (_, _, task) in enumerate(self.tasks):
            if task.task_id == task_id:
                return i + 1
        return None
    
    def get_all_tasks(self) -> List[AITask]:
        """获取所有任务"""
        return [task for _, _, task in self.tasks]
    
    def size(self) -> int:
        """获取队列大小"""
        return len(self.tasks)


class AITaskScheduler:
    """AI任务调度器"""
    
    def __init__(self):
        self.pending_queue = TaskQueue()
        self.running_tasks: Dict[str, AITask] = {}
        self.completed_tasks: Dict[str, AITask] = {}
        self.failed_tasks: Dict[str, AITask] = {}
        
        # 任务处理器注册表
        self.task_processors: Dict[str, Callable] = {}
        
        # 调度配置
        self.max_concurrent_tasks = 10
        self.task_timeout = 3600  # 1小时
        self.cleanup_interval = 300  # 5分钟
        
        # 调度任务
        self._scheduler_task = None
        self._cleanup_task = None
        self._running = False
    
    async def start(self):
        """启动任务调度器"""
        if self._running:
            return
        
        self._running = True
        self._scheduler_task = asyncio.create_task(self._scheduler_loop())
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
        
        logger.info("AI Task Scheduler started")
    
    async def stop(self):
        """停止任务调度器"""
        self._running = False
        
        if self._scheduler_task:
            self._scheduler_task.cancel()
        if self._cleanup_task:
            self._cleanup_task.cancel()
        
        # 取消所有运行中的任务
        for task in self.running_tasks.values():
            task.status = TaskStatus.CANCELLED
            if task.allocation:
                await gpu_resource_manager.release_resource(task.allocation.allocation_id)
        
        logger.info("AI Task Scheduler stopped")
    
    def register_task_processor(self, task_type: str, processor: Callable):
        """注册任务处理器"""
        self.task_processors[task_type] = processor
        logger.info(f"Registered task processor for type: {task_type}")
    
    async def submit_task(self, request: TaskScheduleRequest) -> TaskScheduleResponse:
        """提交任务"""
        try:
            # 创建任务
            task_id = str(uuid4())
            task = AITask(
                task_id=task_id,
                task_type=request.task_type,
                priority=request.priority,
                status=TaskStatus.PENDING,
                resource_request=request.resource_request,
                created_at=datetime.now(),
                metadata=request.metadata
            )
            
            # 添加到队列
            self.pending_queue.add_task(task)
            
            # 获取队列位置
            queue_position = self.pending_queue.get_task_position(task_id)
            
            # 估算开始时间
            estimated_start_time = await self._estimate_start_time(task)
            
            logger.info(f"Task {task_id} submitted with priority {request.priority}")
            
            return TaskScheduleResponse(
                task_id=task_id,
                status=TaskStatus.PENDING,
                estimated_start_time=estimated_start_time,
                queue_position=queue_position
            )
            
        except Exception as e:
            logger.error(f"Failed to submit task: {e}")
            raise
    
    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        try:
            # 检查待处理队列
            if self.pending_queue.remove_task(task_id):
                logger.info(f"Cancelled pending task {task_id}")
                return True
            
            # 检查运行中任务
            if task_id in self.running_tasks:
                task = self.running_tasks[task_id]
                task.status = TaskStatus.CANCELLED
                
                # 释放资源
                if task.allocation:
                    await gpu_resource_manager.release_resource(task.allocation.allocation_id)
                
                # 移动到失败任务列表
                self.failed_tasks[task_id] = task
                del self.running_tasks[task_id]
                
                logger.info(f"Cancelled running task {task_id}")
                return True
            
            return False
            
        except Exception as e:
            logger.error(f"Failed to cancel task {task_id}: {e}")
            return False
    
    async def get_task_status(self, task_id: str) -> Optional[AITask]:
        """获取任务状态"""
        # 检查各个队列
        if task_id in self.running_tasks:
            return self.running_tasks[task_id]
        
        if task_id in self.completed_tasks:
            return self.completed_tasks[task_id]
        
        if task_id in self.failed_tasks:
            return self.failed_tasks[task_id]
        
        # 检查待处理队列
        for task in self.pending_queue.get_all_tasks():
            if task.task_id == task_id:
                return task
        
        return None
    
    async def get_queue_status(self) -> Dict[str, Any]:
        """获取队列状态"""
        return {
            "pending_tasks": self.pending_queue.size(),
            "running_tasks": len(self.running_tasks),
            "completed_tasks": len(self.completed_tasks),
            "failed_tasks": len(self.failed_tasks),
            "max_concurrent_tasks": self.max_concurrent_tasks,
            "registered_processors": list(self.task_processors.keys())
        }
    
    async def get_task_list(self, status: Optional[TaskStatus] = None) -> List[AITask]:
        """获取任务列表"""
        tasks = []
        
        if status is None or status == TaskStatus.PENDING:
            tasks.extend(self.pending_queue.get_all_tasks())
        
        if status is None or status == TaskStatus.RUNNING:
            tasks.extend(self.running_tasks.values())
        
        if status is None or status == TaskStatus.COMPLETED:
            tasks.extend(self.completed_tasks.values())
        
        if status is None or status == TaskStatus.FAILED:
            tasks.extend(self.failed_tasks.values())
        
        return tasks
    
    async def update_task_progress(self, task_id: str, progress: float, metadata: Optional[Dict[str, Any]] = None):
        """更新任务进度"""
        if task_id in self.running_tasks:
            task = self.running_tasks[task_id]
            task.progress = max(0, min(100, progress))
            
            if metadata:
                if task.metadata is None:
                    task.metadata = {}
                task.metadata.update(metadata)
            
            logger.debug(f"Task {task_id} progress updated to {progress}%")
    
    async def _scheduler_loop(self):
        """调度循环"""
        while self._running:
            try:
                await self._process_pending_tasks()
                await asyncio.sleep(1)  # 每秒检查一次
            except Exception as e:
                logger.error(f"Error in scheduler loop: {e}")
                await asyncio.sleep(5)
    
    async def _process_pending_tasks(self):
        """处理待处理任务"""
        # 检查是否可以启动新任务
        if len(self.running_tasks) >= self.max_concurrent_tasks:
            return
        
        # 获取下一个任务
        task = self.pending_queue.get_next_task()
        if not task:
            return
        
        # 检查任务处理器是否存在
        if task.task_type not in self.task_processors:
            logger.error(f"No processor registered for task type: {task.task_type}")
            task.status = TaskStatus.FAILED
            task.error_message = f"No processor registered for task type: {task.task_type}"
            self.failed_tasks[task.task_id] = task
            return
        
        # 尝试分配资源
        allocation = await gpu_resource_manager.allocate_resource(task.resource_request)
        if not allocation:
            # 资源不足，重新加入队列
            self.pending_queue.add_task(task)
            return
        
        # 启动任务
        task.allocation = allocation
        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        self.running_tasks[task.task_id] = task
        
        # 异步执行任务
        asyncio.create_task(self._execute_task(task))
        
        logger.info(f"Started task {task.task_id} on GPU {allocation.gpu_id}")
    
    async def _execute_task(self, task: AITask):
        """执行任务"""
        try:
            processor = self.task_processors[task.task_type]
            
            # 执行任务处理器
            result = await processor(task)
            
            # 任务完成
            task.status = TaskStatus.COMPLETED
            task.completed_at = datetime.now()
            task.progress = 100.0
            
            if isinstance(result, dict):
                if task.metadata is None:
                    task.metadata = {}
                task.metadata.update(result)
            
            # 移动到完成列表
            self.completed_tasks[task.task_id] = task
            del self.running_tasks[task.task_id]
            
            logger.info(f"Task {task.task_id} completed successfully")
            
        except Exception as e:
            # 任务失败
            task.status = TaskStatus.FAILED
            task.completed_at = datetime.now()
            task.error_message = str(e)
            
            # 移动到失败列表
            self.failed_tasks[task.task_id] = task
            del self.running_tasks[task.task_id]
            
            logger.error(f"Task {task.task_id} failed: {e}")
        
        finally:
            # 释放资源
            if task.allocation:
                await gpu_resource_manager.release_resource(task.allocation.allocation_id)
    
    async def _cleanup_loop(self):
        """清理循环"""
        while self._running:
            try:
                await self._cleanup_old_tasks()
                await self._check_task_timeouts()
                await asyncio.sleep(self.cleanup_interval)
            except Exception as e:
                logger.error(f"Error in cleanup loop: {e}")
                await asyncio.sleep(60)
    
    async def _cleanup_old_tasks(self):
        """清理旧任务"""
        cutoff_time = datetime.now() - timedelta(hours=24)
        
        # 清理完成的任务
        completed_to_remove = [
            task_id for task_id, task in self.completed_tasks.items()
            if task.completed_at and task.completed_at < cutoff_time
        ]
        
        for task_id in completed_to_remove:
            del self.completed_tasks[task_id]
        
        # 清理失败的任务
        failed_to_remove = [
            task_id for task_id, task in self.failed_tasks.items()
            if task.completed_at and task.completed_at < cutoff_time
        ]
        
        for task_id in failed_to_remove:
            del self.failed_tasks[task_id]
        
        if completed_to_remove or failed_to_remove:
            logger.info(f"Cleaned up {len(completed_to_remove)} completed and {len(failed_to_remove)} failed tasks")
    
    async def _check_task_timeouts(self):
        """检查任务超时"""
        current_time = datetime.now()
        timeout_threshold = timedelta(seconds=self.task_timeout)
        
        timed_out_tasks = []
        
        for task_id, task in self.running_tasks.items():
            if task.started_at and (current_time - task.started_at) > timeout_threshold:
                timed_out_tasks.append(task_id)
        
        for task_id in timed_out_tasks:
            task = self.running_tasks[task_id]
            task.status = TaskStatus.FAILED
            task.completed_at = current_time
            task.error_message = "Task timeout"
            
            # 释放资源
            if task.allocation:
                await gpu_resource_manager.release_resource(task.allocation.allocation_id)
            
            # 移动到失败列表
            self.failed_tasks[task_id] = task
            del self.running_tasks[task_id]
            
            logger.warning(f"Task {task_id} timed out after {self.task_timeout} seconds")
    
    async def _estimate_start_time(self, task: AITask) -> Optional[datetime]:
        """估算任务开始时间"""
        try:
            # 简单估算：基于队列位置和平均任务执行时间
            queue_position = self.pending_queue.get_task_position(task.task_id)
            if queue_position is None:
                return None
            
            # 假设平均任务执行时间为10分钟
            average_task_duration = 600  # 秒
            
            # 考虑当前运行的任务数量
            running_tasks_count = len(self.running_tasks)
            available_slots = max(0, self.max_concurrent_tasks - running_tasks_count)
            
            if available_slots > 0 and queue_position <= available_slots:
                # 可以立即开始
                return datetime.now()
            else:
                # 需要等待
                wait_time = ((queue_position - available_slots) / self.max_concurrent_tasks) * average_task_duration
                return datetime.now() + timedelta(seconds=wait_time)
                
        except Exception as e:
            logger.error(f"Failed to estimate start time: {e}")
            return None
    
    def set_max_concurrent_tasks(self, max_tasks: int):
        """设置最大并发任务数"""
        self.max_concurrent_tasks = max(1, max_tasks)
        logger.info(f"Max concurrent tasks set to {self.max_concurrent_tasks}")
    
    def set_task_timeout(self, timeout_seconds: int):
        """设置任务超时时间"""
        self.task_timeout = max(60, timeout_seconds)  # 最少1分钟
        logger.info(f"Task timeout set to {self.task_timeout} seconds")


# 全局任务调度器实例
ai_task_scheduler = AITaskScheduler()